#include <iostream>
#include <string>
#include <iomanip>
#include "display.h"
#include "utils.h"
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <signal.h>
#include <sys/wait.h>
#include "opencv2/core/core.hpp"
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include "net.h"
#include <benchmark.h>
#include "SocketMatTransmissionServer.h"


using namespace std;
using namespace cv;

#define WIDTH 224
#define HEIGHT 224
#define BUF_SIZE 100

void read_childproc(int sig);

// 对输入图片进行分类, 获取预测结果, pred向量中各类别置信度的排列是无序的
// TODO: 目前是模拟的结果, 分类可用之后, 使用真实分类流程替代
float *predict()
{
    float *pred = new float[NUM_CLASSES];
    float small[7] = {0.01, 0.0, 0.01, 0.01, 0.02, 0.03, 0.02};

    srand((unsigned) time(NULL));
    float top1 = (rand() % 10) / 10.f;
    float top2 = 0.9f - top1;

    srand((unsigned) time(NULL));
    int index_top1 = rand() % 9;
    int index_top2 = index_top1 / 2 == index_top1 ? index_top1 + 1 : index_top1 / 2;

    int idx_small = 0;
    int i;
    for (i = 0; i < NUM_CLASSES; i++)
    {
        pred[i] = i == index_top1 ? top1
                                  : i == index_top2 ? top2
                                                    : small[idx_small++];
    }

    return pred;
}

float * detect_mobilenetv2(const cv::Mat &bgr, ncnn::Net &mobilenetv2)
{
    // ncnn::Net mobilenetv2;

    // mobilenetv2.opt.use_vulkan_compute = true;
    // mobilenetv2.opt.num_threads = 8;

    // mobilenetv2.opt.use_winograd_convolution = true;
    // mobilenetv2.opt.use_sgemm_convolution = true;
    // mobilenetv2.opt.use_fp16_packed = true;
    // mobilenetv2.opt.use_fp16_storage = true;
    // mobilenetv2.opt.use_fp16_arithmetic = true;
    // mobilenetv2.opt.use_packing_layout = true;
    // mobilenetv2.opt.use_shader_pack8 = false;
    // mobilenetv2.opt.use_image_storage = false;

     float *pred = new float[NUM_CLASSES];

    // // https://github.com/miaow1988/ShuffleNet_V2_pytorch_caffe
    // // models can be downloaded from https://github.com/miaow1988/ShuffleNet_V2_pytorch_caffe/releases
    // mobilenetv2.load_param("../model/mobilenetv2-opt.param");
    // mobilenetv2.load_model("../model/mobilenetv2-opt.bin");

    ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, 224, 224);


    const float mean_vals[3] = {0.485f * 255.f, 0.456f * 255.f, 0.406f * 255.f};
    const float norm_vals[3] = {1 / 0.229f / 255.f, 1 / 0.224f / 255.f, 1 / 0.225f / 255.f};
    in.substract_mean_normalize(mean_vals, norm_vals);

    ncnn::Extractor ex = mobilenetv2.create_extractor();

    ex.input("input", in);

    ncnn::Mat out;
    ex.extract("output", out);

    // manually call softmax on the fc output
    // convert result into probability
    // skip if your model already has softmax operation
    {
        ncnn::Layer *softmax = ncnn::create_layer("Softmax");

        ncnn::ParamDict pd;
        softmax->load_param(pd);

        softmax->forward_inplace(out, mobilenetv2.opt);

        delete softmax;
    }

    out = out.reshape(out.w * out.h * out.c);


    for (int j = 0; j < out.w && j< NUM_CLASSES; j++)
    {
        pred[j] = out[j];
    }

    return pred;
}

int maxArray(float *pred)
{
    float max=0;
    int index=0;
    for(int i=0;i<NUM_CLASSES;i++)
    {

        if(max< pred[i]){
            max= pred[i]; 
            index = i;
        }      
        
    }
    return index;
}

// 对视频进行分类并展示
int run_predict(Mat frame, ncnn::Net &mobilenetv2)
{
    resize(frame, frame, Size(640, 480));

    float *pred = detect_mobilenetv2(frame, mobilenetv2);
    //show_result(frame, pred, TOPK, true, false);

    //printf("1:%f, 2:%f,3:%f,4:%f,5:%f,6:%f,7:%f,8:%f,9:%f",
    //    pred[0],pred[1],pred[2],pred[3],pred[4],pred[5],pred[6],pred[7],pred[8]);

    int res = maxArray(pred);
    delete[] pred;
    return res;
      
}




int main(int argc, char **argv)
{
    int i;
    pid_t pid;
    SocketMatTransmissionServer server;
	struct sigaction act;
    for (i = 0; i < argc; ++i)
    {
        if (!argv[i]) continue;
        strip_args(argv[i]);
    }

    if (argc < 2)
    {
        fprintf(stderr, "usage: %s <port>\n", argv[0]);
        return 0;
    }
	act.sa_handler=read_childproc;
	sigemptyset(&act.sa_mask);
	act.sa_flags=0;
	state=sigaction(SIGCHLD,&act,0);

    if(server.socketBind(atoi(argv[1])) < 0)
    {
        fprintf(stderr, "port %s bind fail\n", argv[1]);
        return 0;
    }

    

    while(1)
    {
        server.socketAccept();
        if(server.sockConn ==-1) 
            continue;
       

        pid =fork();
        if (pid==-1)
        {
            close(server.sockConn);
            continue;
        }
        if(pid ==0){
            close(server.server_sockfd);

            Mat image;
            int res = 0;

            // 接受帧信息，比如帧的高，宽
            struct infohead head;
            res = server.receive(head);
            server.detectType=head.detectType;
            server.width=head.width;
            server.height=head.height;
            server.packageNum=head.packageNum;

            // 设置Net对象的参数
            ncnn::Net mobilenetv2;
            mobilenetv2.opt.use_vulkan_compute = true;
            mobilenetv2.opt.num_threads = 8;
            mobilenetv2.opt.use_winograd_convolution = true;
            mobilenetv2.opt.use_sgemm_convolution = true;
            mobilenetv2.opt.use_fp16_packed = true;
            mobilenetv2.opt.use_fp16_storage = true;
            mobilenetv2.opt.use_fp16_arithmetic = true;
            mobilenetv2.opt.use_packing_layout = true;
            mobilenetv2.opt.use_shader_pack8 = false;
            mobilenetv2.opt.use_image_storage = false; 

            // 加载模型
            mobilenetv2.load_param("../model/mobilenetv2-opt.param");
            mobilenetv2.load_model("../model/mobilenetv2-opt.bin");


            // 接受图像帧
            while (1)
            {
                res = server.receive(image);
                if(res > 0)
                {
                    server.send_res(run_predict(image,mobilenetv2));

                    if(argc > 2)
                        server.writeOpencvMat(image);
                }
                else
                {
                    break;
                }
            }       
            server.socketDisconnect();
        }
        else
        {

            close(server.sockConn);
            continue;
        }
        break;

            

   }

    return 0;
}

void read_childproc(int sig){

	pid_t pid;
	int status ;
	pid =waitpid(-1,&status,WNOHANG);
	printf("removed proc id: %d\n",pid);
}
